#coding: utf-8
from capstone import *
from capstone.x86 import *
from ctypes import wintypes
import os

def ReadMemory(addr, numbytes):
	bytes_read = ctypes.create_string_buffer(numbytes)
	ret = ctypes.c_size_t()
	rRd(hprocess, addr, bytes_read, numbytes, ctypes.byref(ret))
	r = []
	for i in bytes_read.raw:
		r.append(int(i))
	return bytes(r)

def IsBranch(s):
	br = ['loopnz', 'loopne', 'loopz', 'loope', 'loop', 'ja', 'jae', 'jb', 'jbe', 'jc', 'jcxz', 'jecxz', 'je', 'jg', 'jge', 'jl', 'jle', 'jna', 'jnae', 'jnb', 'jnbe', 'jnc', 'jne', 'jng', 'jnge', 'jnl', 'jnle', 'jno', 'jnp', 'jns', 'jnz', 'jo', 'jp', 'jpe', 'jpo', 'js', 'jz', 'jmp']
	return (s in br)

#获取一段汇编代码，以RET X指令作为终止信号
#CALL指令转换为push ret/jmp addr
#条件跳转默认成立
def GetAsm(StartAddr):
	LstAsm = []
	ObDis = Cs(CS_ARCH_X86, CS_MODE_32)
	ObDis.detail = True
	k = b'\x6A\x00'
	asmm = ObDis.disasm(k, 0x401000)
	for insn in asmm:
		vcall = insn
	state = True
	while state:
		code = ReadMemory(StartAddr, 30)
		asmm = ObDis.disasm(code, StartAddr)
		x = 0
		for insn in asmm:
			if insn.mnemonic == 'ret':
				LstAsm.append(insn)
				state = False
			if x == 0:
				if IsBranch(insn.mnemonic):
					for i in insn.operands:
						StartAddr = i.imm
					break
				elif insn.mnemonic == 'call':
					for i in insn.operands:
						LstAsm.append(vcall)
						StartAddr = i.imm
					break
				else:
					LstAsm.append(insn)
			elif x == 1:
				StartAddr = insn.address
				break
			x += 1
	return LstAsm        

def ReadReg(RegName, Insninfo):
	(regs_read, regs_write) = Insninfo.regs_access()
	for r in regs_read:
		if Insninfo.reg_name(r) == RegName:
			return True
	return False	

def ModifyReg(RegName, Insninfo):
	(regs_read, regs_write) = Insninfo.regs_access()
	for r in regs_write:
		if Insninfo.reg_name(r) == RegName:
			return True
	return False

def ModifyRegExact(RegName, Insninfo):
	DicReg = {}
	DicReg['eax'] = []
	DicReg['eax'].append('eax')
	DicReg['eax'].append('ax')
	DicReg['eax'].append('al')
	DicReg['eax'].append('ah')
	DicReg['ebx'] = []
	DicReg['ebx'].append('ebx')
	DicReg['ebx'].append('bx')
	DicReg['ebx'].append('bl')
	DicReg['ebx'].append('bh')	
	DicReg['ecx'] = []
	DicReg['ecx'].append('ecx')
	DicReg['ecx'].append('cx')
	DicReg['ecx'].append('cl')
	DicReg['ecx'].append('ch')
	DicReg['edx'] = []
	DicReg['edx'].append('edx')
	DicReg['edx'].append('dx')
	DicReg['edx'].append('dl')
	DicReg['edx'].append('dh')
	DicReg['ebp'] = []
	DicReg['ebp'].append('ebp')
	DicReg['ebp'].append('bp')
	DicReg['edi'] = []
	DicReg['edi'].append('edi')
	DicReg['edi'].append('di')
	DicReg['esi'] = []
	DicReg['esi'].append('esi')
	DicReg['esi'].append('si')
	for i in DicReg[RegName]:
		if ModifyReg(i, Insninfo):
			return True
	return False

def ReadRegExact(RegName, Insninfo):
	DicReg = {}
	DicReg['eax'] = []
	DicReg['eax'].append('eax')
	DicReg['eax'].append('ax')
	DicReg['eax'].append('al')
	DicReg['eax'].append('ah')
	DicReg['ebx'] = []
	DicReg['ebx'].append('ebx')
	DicReg['ebx'].append('bx')
	DicReg['ebx'].append('bl')
	DicReg['ebx'].append('bh')	
	DicReg['ecx'] = []
	DicReg['ecx'].append('ecx')
	DicReg['ecx'].append('cx')
	DicReg['ecx'].append('cl')
	DicReg['ecx'].append('ch')
	DicReg['edx'] = []
	DicReg['edx'].append('edx')
	DicReg['edx'].append('dx')
	DicReg['edx'].append('dl')
	DicReg['edx'].append('dh')
	DicReg['ebp'] = []
	DicReg['ebp'].append('ebp')
	DicReg['ebp'].append('bp')
	DicReg['edi'] = []
	DicReg['edi'].append('edi')
	DicReg['edi'].append('di')
	DicReg['esi'] = []
	DicReg['esi'].append('esi')
	DicReg['esi'].append('si')
	for i in DicReg[RegName]:
		if ReadReg(i, Insninfo):
			return True
	return False

def AnalyzeStackObfs(LstAsm):
	global MaxPtr
	ptr = -1 #模拟栈指针，但是从低地址向高地址扩展的，以DWORD为最小单位，即执行一次push dword栈指针递增1
	LstStackInfo = []
	stackele = {}
	stackele['read'] = []
	stackele['write'] = []
	for insn in LstAsm:
		stackele = {}
		stackele['read'] = []
		stackele['write'] = []		
		if insn.mnemonic == 'push': #入栈
			for x in insn.operands:
				if x.type == X86_OP_MEM and insn.reg_name(x.mem.base) == 'esp' and x.mem.index == 0:
					if x.mem.disp % 4 != 0:
						print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
					if x.access == CS_AC_READ:
						stackele['read'].append(ptr - x.mem.disp // 4)
					elif x.access == CS_AC_WRITE:
						stackele['write'].append(ptr - x.mem.disp // 4)
					elif x.access == CS_AC_READ | CS_AC_WRITE:
						stackele['read'].append(ptr - x.mem.disp // 4)
						stackele['write'].append(ptr - x.mem.disp // 4)		
			ptr += 1
			stackele['write'].append(ptr)
		elif insn.mnemonic in ['pushad', 'pushal']: #入栈
			ptr += 8
			for x in range(ptr, ptr - 8, -1):
				stackele['write'].append(x)
		elif insn.mnemonic == 'pushfd': #入栈
			ptr += 1
			stackele['write'].append(ptr)
		elif insn.mnemonic == 'pop': #出栈
			for x in insn.operands:
				if x.type == X86_OP_MEM and insn.reg_name(x.mem.base) == 'esp' and x.mem.index == 0:
					if x.mem.disp % 4 != 0:
						print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
					if x.access == CS_AC_READ:
						stackele['read'].append(ptr - x.mem.disp // 4)
					elif x.access == CS_AC_WRITE:
						stackele['write'].append(ptr - x.mem.disp // 4)
					elif x.access == CS_AC_READ | CS_AC_WRITE:
						stackele['read'].append(ptr - x.mem.disp // 4)
						stackele['write'].append(ptr - x.mem.disp // 4)
			stackele['read'].append(ptr)
			ptr -= 1
		elif insn.mnemonic == 'popfd': #出栈
			stackele['read'].append(ptr)
			ptr -= 1
		elif insn.mnemonic in ['popad', 'popal']: #出栈
			for x in range(ptr, ptr - 8, -1):
				stackele['read'].append(x)
			ptr -= 8
		elif insn.mnemonic in ['sub', 'sbb'] and ModifyReg('esp', insn): #调整栈指针
			for x, k in zip(insn.operands, range(1, 3)):
				if k == 2:
					if x.type == X86_OP_IMM:
						if x.imm % 4 == 0:
							ptr += (x.imm // 4)
						else:
							print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
					else:
						print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
		elif insn.mnemonic in ['add', 'adc'] and ModifyReg('esp', insn): #调整栈指针
			for x, k in zip(insn.operands, range(1, 3)):
				if k == 2:
					if x.type == X86_OP_IMM:
						if x.imm % 4 == 0:
							ptr -= (x.imm // 4)
						else:
							print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
					else:
						print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))			
		elif insn.mnemonic == 'lea': #调整栈指针
			if ModifyReg('esp', insn):
				for x, k in zip(insn.operands, range(1, 3)):
					if k == 2:
						if x.type == X86_OP_MEM:
							if x.mem.disp % 4 == 0:
								ptr -= (x.mem.disp // 4)
							else:
								print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
						else:
							print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
		elif insn.mnemonic == 'ret':
			stackele['read'].append(ptr)
			for x in insn.operands:
				ptr -= (x.imm // 4 + 1)
		elif ModifyReg('esp', insn): #如果还有其它调整栈指针的指令没有被考虑，打印出详细信息
			print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))
		else: #处理 xxx [esp + xx],xx格式指令
			for x in insn.operands:
				if x.type == X86_OP_MEM and insn.reg_name(x.mem.base) == 'esp' and x.mem.index == 0:
					if x.mem.disp % 4 != 0:
						print('栈混淆清理可能有误，因为有未知的调整栈指针指令，该指令位于%X' % (insn.address))

					if x.access == CS_AC_READ:
						stackele['read'].append(ptr - x.mem.disp // 4)
					elif x.access == CS_AC_WRITE:
						stackele['write'].append(ptr - x.mem.disp // 4)
					elif x.access == CS_AC_READ | CS_AC_WRITE:
						stackele['read'].append(ptr - x.mem.disp // 4)
						stackele['write'].append(ptr - x.mem.disp // 4)	
		LstStackInfo.append(stackele)
		if ptr > MaxPtr:
			MaxPtr = ptr

	return LstStackInfo

#判断b是否包含a
def IsListIn(a, b):
	for i in a:
		if i not in b:
			return False
	return True

def ListFind(lst, a):
	i = 0
	for x in lst:
		if a == x:
			return i
		i += 1
	return -1

def CleanStackObfs(LstStackInfo, LstAsm):
	def AnalyzeStackVar(LstStackInfo, StackVar):
		k = True
		LstStackInfo.reverse()
		for info in LstStackInfo:
			info[StackVar] = k
			if StackVar in info['write']:
				k = False
			if StackVar in info['read']:
				k = True
		LstStackInfo.reverse()
		return
	def GetUnVar(info):
		z = []
		for x in range(0, MaxPtr + 1):
			if not info[x]:
				z.append(x)

		return z
	#默认在尾部时，所有栈变量都是活跃的
	for i in range(0, MaxPtr + 1):
		AnalyzeStackVar(LstStackInfo, i)

	LstNewAsm = LstAsm[:]
	LstNewStackInfo = LstStackInfo[:]
	LstAsmInfo = []
	for x, k in zip(LstNewAsm, LstNewStackInfo):
		y = {}
		y['Asm'] = x
		y['Stack'] = k
		LstAsmInfo.append(y)
	LstNewAsmInfo = LstAsmInfo[:]
	ObDis = Cs(CS_ARCH_X86, CS_MODE_32)
	ObDis.detail = True
	for insn in ObDis.disasm(b'\x8D\x64\x24\xFC', 401000):
		vPush = insn
	for insn in ObDis.disasm(b'\x8D\x64\x24\xE0', 401000):
		vPushal = insn
	for insn in ObDis.disasm(b'\x8D\x64\x24\x04', 401000):
		vPop = insn
	for insn in ObDis.disasm(b'\x8D\x64\x24\x20', 401000):
		vPopal = insn
	for asmi in LstAsmInfo:
		if (len(asmi['Stack']['write']) != 0) and (IsListIn(asmi['Stack']['write'], GetUnVar(asmi['Stack']))):
			if asmi['Asm'].mnemonic not in ['push', 'pushal', 'pushfd', 'pushad', 'pop', 'popfd', 'popal', 'popad']: #不涉及到栈指针调整的指令可以直接删除
				print('位于%X处的指令被清除了' % asmi['Asm'].address)
				LstNewAsmInfo.remove(asmi)
			if asmi['Asm'].mnemonic in ['push', 'pushfd']:
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Asm'] = vPush
				print('位于%X处的指令被清除了' % asmi['Asm'].address)
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['write'] = []
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['read'] = []
			if asmi['Asm'].mnemonic in ['pushal', 'pushad']:
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Asm'] = vPushal
				print('位于%X处的指令被清除了' % asmi['Asm'].address)
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['write'] = []
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['read'] = []	
			if asmi['Asm'].mnemonic in ['pop', 'popfd']:
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Asm'] = vPop
				print('位于%X处的指令被清除了' % asmi['Asm'].address)
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['write'] = []
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['read'] = []
			if asmi['Asm'].mnemonic in ['popad', 'popal']:
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Asm'] = vPopal
				print('位于%X处的指令被清除了' % asmi['Asm'].address)
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['write'] = []
				LstNewAsmInfo[ListFind(LstNewAsmInfo, asmi)]['Stack']['read'] = []				

	LstAsmInfo = LstNewAsmInfo[:]

	return LstNewAsmInfo
def CleanRegObfs(LstAsm):
	def IsUseful(AsmInfo):
		insn = AsmInfo['Asm']
		(regs_read, regs_write) = insn.regs_access()
		if insn.mnemonic == 'pop':
			for r in insn.operands:
				if r.type == X86_OP_REG:
					AsmInfo[insn.reg_name(r.reg)] = False
					return AsmInfo[insn.reg_name(r.reg)]
				else:
					return True
		if len(regs_write) == 0:
			return True
		if ModifyReg('esp', insn):
			return True
		for r in AsmInfo.keys():
			if r != 'Asm' and AsmInfo[r]: #如果影响了任意一个活跃寄存器即视为有效
				if ModifyRegExact(r, insn):
					return True 
		return False
	def AnalyzeReg(LstAsm, 	Reg):
		k = True
		LstAsm.reverse()
		LstRegInfo = []
		for insn in LstAsm:
			LstRegInfo.insert(0, k)
			if ModifyRegExact(Reg, insn):
				k = False
			if ReadRegExact(Reg, insn):
				k = True
		LstAsm.reverse()
		return LstRegInfo
	LstAsmInfo = []
	LstReg = ['eax', 'ebx', 'ecx', 'edx', 'edi', 'esi', 'ebp']
	LstRegInfo = {}
	for r in LstReg:
		LstRegInfo[r] = AnalyzeReg(LstAsm, r)
	for i in range(0, len(LstAsm)):
		insn = LstAsm[i]
		asmi = {}
		asmi['Asm'] = insn
		for k in LstRegInfo.keys():
			asmi[k] = LstRegInfo[k][i]
		LstAsmInfo.append(asmi)

	LstNewAsmInfo = LstAsmInfo[:]
	LstDel = []
	ObDis = Cs(CS_ARCH_X86, CS_MODE_32)
	ObDis.detail = True
	for insn in ObDis.disasm(b'\x8D\x64\x24\xFC', 401000):
		vPush = insn
	for insn in ObDis.disasm(b'\x8D\x64\x24\xE0', 401000):
		vPushal = insn
	for insn in ObDis.disasm(b'\x8D\x64\x24\x04', 401000):
		vPop = insn
	for insn in ObDis.disasm(b'\x8D\x64\x24\x20', 401000):
		vPopal = insn
	for i in LstAsmInfo:
		if not IsUseful(i):
			if i['Asm'].mnemonic == 'pop':
				LstNewAsmInfo[ListFind(LstNewAsmInfo, i)]['Asm'] = vPop
			else:
				print('%X' % i['Asm'].address)
				LstDel.append(i)
		
	for i in LstDel:
		LstNewAsmInfo.remove(i)

	LstAsmInfo = LstNewAsmInfo[:]
	return LstAsmInfo

	

def DisDispatch():
	LstAsm = GetAsm(start)
	ObLog.write('Dispatch(未经过处理):\n')
	for i in LstAsm:
		ObLog.write(i.mnemonic + '	' + i.op_str + '\n')
	LstAsmInfo = CleanStackObfs(AnalyzeStackObfs(LstAsm), LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstAsmInfo = CleanRegObfs(LstAsm)
	LstAsm = []	
	for i in LstAsmInfo:
		LstAsm.append(i['Asm'])
	LstStackInfo = AnalyzeStackObfs(LstAsm)
	for i, x in zip(LstAsm, LstStackInfo):
		print('%X %s %s' % (i.address, i.mnemonic, i.op_str), 'read =', x['read'], 'write =', x['write'])
	return

def Main():
	print('Fuck VMProtect V2.X')
	print('开始分析Dispatch...')
	DisDispatch()
	rOc = ctypes.WinDLL('kernel32.dll').CloseHandle
	rOc.argtype = [wintypes.DWORD]
	rOc.restype = wintypes.BOOL
	rOc(hprocess)
	return

first = False
MaxPtr = 0
ObLog = open('analyse.txt', 'w')
pid = int(input('请输入进程ID: '))
start = int(input('请输入VM代码开始地址(HEX): '), 16)
rOp = ctypes.WinDLL('kernel32.dll').OpenProcess
rOp.argtypes = [wintypes.DWORD, wintypes.DWORD, wintypes.DWORD]
rOp.restype = wintypes.DWORD
hprocess = rOp(2035711, 0, pid)
rRd = ctypes.WinDLL('kernel32.dll').ReadProcessMemory
rRd.argtypes = [wintypes.HANDLE,wintypes.LPCVOID,wintypes.LPVOID,ctypes.c_size_t,ctypes.POINTER(ctypes.c_size_t)]
Main()
ObLog.close()